import numpy as np

# from environ.env_5bits_2048 import Env_2048

# env = Env_2048()


class Tuple_17_4():
    def __init__(self) -> None:
        self.table = np.zeros((17, 17, 17, 17, 17), dtype=np.float64)
        self.index_table = np.zeros((17, 5), dtype=np.int32)
        self.index_table[:, 0] = np.arange(17)

    def get_value(self, board:np.array):
        index_0 = board[0, :]
        index_1 = board[1, :]
        index_2 = board[2, :]
        index_3 = board[3, :]
        index_4 = board[:, 0]
        index_5 = board[:, 1]
        index_6 = board[:, 2]
        index_7 = board[:, 3]
        index_8 = board[np.array([0, 0, 1, 1]), np.array([0, 1, 1, 0])]
        index_9 = board[np.array([0, 0, 1, 1]), np.array([1, 2, 2, 1])]
        index_10= board[np.array([0, 0, 1, 1]), np.array([2, 3, 3, 2])]
        index_11= board[np.array([1, 1, 2, 2]), np.array([0, 1, 1, 0])]
        index_12= board[np.array([1, 1, 2, 2]), np.array([1, 2, 2, 1])]
        index_13= board[np.array([1, 1, 2, 2]), np.array([2, 3, 3, 2])]
        index_14= board[np.array([2, 2, 3, 3]), np.array([0, 1, 1, 0])]
        index_15= board[np.array([2, 2, 3, 3]), np.array([1, 2, 2, 1])]
        index_16= board[np.array([2, 2, 3, 3]), np.array([2, 3, 3, 2])]

        self.index_table[0, 1:] = index_0
        self.index_table[1, 1:] = index_1
        self.index_table[2, 1:] = index_2
        self.index_table[3, 1:] = index_3
        self.index_table[4, 1:] = index_4
        self.index_table[5, 1:] = index_5
        self.index_table[6, 1:] = index_6
        self.index_table[7, 1:] = index_7
        self.index_table[8, 1:] = index_8
        self.index_table[9, 1:] = index_9
        self.index_table[10, 1:] = index_10
        self.index_table[11, 1:] = index_11
        self.index_table[12, 1:] = index_12
        self.index_table[13, 1:] = index_13
        self.index_table[14, 1:] = index_14
        self.index_table[15, 1:] = index_15
        self.index_table[16, 1:] = index_16

        return np.sum(
        self.table[self.index_table[:,0], 
                    self.index_table[:,1], 
                    self.index_table[:,2], 
                    self.index_table[:,3], 
                    self.index_table[:,4]]
        )

    def update_value(self, delta:np.float32):
        self.table[self.index_table[:,0], 
                    self.index_table[:,1], 
                    self.index_table[:,2], 
                    self.index_table[:,3], 
                    self.index_table[:,4]] -= delta
 


if __name__ == '__main__':
    t = Tuple_17_4()

    board=np.array([[0, 0, 0, 0],
                            [0, 0, 0, 0],
                            [0, 0, 0, 0], 
                            [0, 0, 0, 0]], dtype=np.uint16)
    
    print(
    t.get_value(board)
    )

    t.update_value(0.01)

    print(np.where(t.table!=0))




